#!/usr/bin/env python3
'''
正在适配

参考链接：https://github.com/lutery/twm.git

训练记录：
'''
import gymnasium as gym
import argparse

import torch
import torch.multiprocessing as mp

import yaml
import pathlib
import sys

from lib import model, common, trainer

import ale_py

gym.register_envs(ale_py)


if __name__ == "__main__":
    mp.set_start_method('spawn')
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action='store_true', help='Enable CUDA')
    parser.add_argument('--configs', nargs='+', default=['defaults'])
    parser.add_argument("-n", "--name", default='laserGates', help="Name of the run")
    args, remaining = parser.parse_known_args()
    device = common.select_device(args=args)
    if 'cuda' in device.type:
        torch.backends.cudnn.benchmark = True 

     # parser.add_argument('--configs', nargs='+', required=True)
    # Comment the line above and comment out the line below if you want to debug in IDE like PyCharm
    # Update from configs.yaml
    configs = yaml.safe_load((pathlib.Path(sys.argv[0]).parent / 'config/config.yaml').read_text(encoding='utf-8'))
    default_params = dict()
    for name in args.configs:
        default_params.update(configs[name])
    # Update from cli
    for key, value in default_params.items():
        parser.add_argument('--' + key, type=type(value), default=value)
    args = parser.parse_args(remaining)
    params = vars(args)
    params['env_reward_transform'] = model.get_activation(params['env_reward_transform'])
    params['obs_act'] = model.get_activation(params['obs_act'])
    params['dyn_act'] = model.get_activation(params['dyn_act'])
    params['ac_act'] = model.get_activation(params['ac_act'])
    

    trainer = trainer.Trainer(params=params, device=device)
    trainer.load_model()
    trainer.train_model()

